#ifndef TENSOR_H
#define TENSOR_H

#include <cstddef>
#include <cstdint>
#include <iostream>
#include <vector>
#include <numeric>

// 定义数据类型枚举，用于表示张量支持的数据类型
enum class DataType {
    FLOAT32,
    INT32
};

enum class DataLayout {
    HWC,
    CHW
};

struct Tensor {
    DataLayout layout;
    void* data; 
    int num_dims;
    std::vector<int> dims;  // 存储各维度大小的数组，如{batch_size, channels, height, width}
    DataType data_type;  // 数据类型，比如 float、double 等
    bool is_gpu;
};

// 获取对应数据类型的字节大小
size_t getDataTypeSize(DataType data_type);

// 在 CPU 上分配内存
Tensor allocateTensorOnCPU(std::vector<int> dims, DataType data_type, DataLayout layout = DataLayout::HWC);

// 在 GPU 上分配内存
Tensor allocateTensorOnGPU(std::vector<int> dims, DataType data_type, DataLayout layout = DataLayout::HWC);

// 释放张量内存（根据是否在 GPU 上选择相应释放方式）
void freeTensor(Tensor tensor);

// 在 CPU 和 GPU 之间拷贝数据（例如从 CPU 到 GPU）
void copyToGPU(Tensor src, Tensor dst);

#endif // TENSOR_H